############################################################
#   File: engine.py                                        #
#   Created: 2019-11-20 15:02:13                           #
#   Author : wvinzh                                        #
#   Email : wvinzh@qq.com                                  #
#   ------------------------------------------             #
#   Description:engine.py                                  #
#   Copyright@2019 wvinzh, HUST                            #
############################################################
import time
from attention import calculate_pooling_center_loss
from attention import attention_crop
from config import getDatasetConfig, getConfig
from myutils import getLogger
from myutils import accuracy, get_lr, save_checkpoint, set_seed
from meter import AverageMeter
from rank_loss import RankLoss, MaxMinNormal

import torch
import torch.nn.functional as F

class Engine():
    def __init__(self,):
        pass

    def train(self,state,epoch):
        batch_time = AverageMeter()
        data_time = AverageMeter()
        losses = AverageMeter()
        top1 = AverageMeter()
        top5 = AverageMeter()
        rk = RankLoss()
        config = state['config']
        print_freq = config.print_freq
        model = state['model']
        criterion = state['criterion']
        optimizer = state['optimizer']
        train_loader = state['train_loader']
        model.train()
        end = time.time()

        for i, (img, label) in enumerate(train_loader):
            # measure data loading time
            data_time.update(time.time() - end)

            target = label.cuda()
            input = img.cuda()
            # compute output
            attention_maps, raw_features, output1 = model(input)
            features = raw_features.reshape(raw_features.shape[0], -1)

            feature_center_loss, center_diff = calculate_pooling_center_loss(
                features, state['center'], target, alfa=config.alpha)

            # update model.centers
            state['center'][target] += center_diff

            # compute refined loss
            # img_drop = attention_drop(attention_maps,input)
            # img_crop = attention_crop(attention_maps, input)
            ##################################################
            ###### crop the image ############################
            ##################################################
            images, topks = attention_crop(attention_maps, input)
            b, k, c, w, h = images.shape
            img_crop = images.view(b*k, c, w, h)
            _, _, output2 = model(img_crop)
            target2 = target.repeat_interleave(4)
            loss1 = criterion(output1, target)
            loss2 = criterion(output2, target2)

            # topk : b*4
            # ouput2 : (b*4)*200
            # target2: (b*4)
            # pp: (b*4)*1
            ##################################################
            ######## rank loss ###############################
            ##################################################
            topks = MaxMinNormal()(topks)
            pp = output2.gather(-1, target2.unsqueeze(1))
            pp = pp.view(b, -1)
            rank_loss = rk(topks, pp)


            loss = (loss1+loss2)/2 + feature_center_loss + rank_loss
            # measure accuracy and record loss
            prec1, prec5 = accuracy(output1, target, topk=(1, 5))
            losses.update(loss.item(), input.size(0))
            top1.update(prec1[0], input.size(0))
            top5.update(prec5[0], input.size(0))

            # compute gradient and do SGD step
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()
        return top1.avg, losses.avg
    
    def validate(self,state):
        batch_time = AverageMeter()
        losses = AverageMeter()
        top1 = AverageMeter()
        top5 = AverageMeter()
        
        config = state['config']
        print_freq = config.print_freq
        model = state['model']
        val_loader = state['val_loader']
        criterion = state['criterion']
        # switch to evaluate mode
        model.eval()
        with torch.no_grad():
            end = time.time()
            for i, (input, target) in enumerate(val_loader):
                target = target.cuda()
                input = input.cuda()
                # forward
                attention_maps, raw_features, output1 = model(input)
                features = raw_features.reshape(raw_features.shape[0], -1)
                feature_center_loss, _ = calculate_pooling_center_loss(
                    features, state['center'], target, alfa=config.alpha)

                images, topks = attention_crop(attention_maps, input)
                b, k, c, w, h = images.shape
                img_crop = images.view(b * k, c, w, h)
                _, _, output2 = model(img_crop)

                bb, num_class = output2.shape
                output2 = output2.view(b, k, -1)
                output2 = torch.nn.functional.softmax(output2, -1)
                output2 = output2.mean(1)

                output1 = torch.nn.functional.softmax(output1, -1)
                output2 = (output1 + output2)/2

                prec1, prec5 = accuracy(output2, target, topk=(1, 5))
                top1.update(prec1[0], input.size(0))
                top5.update(prec5[0], input.size(0))

            print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
                .format(top1=top1, top5=top5))

        return top1.avg, top5.avg

if __name__ == '__main__':

    engine = Engine()
    engine.train()

